package cn.doitedu.ml.util

import org.apache.spark.ml.linalg.Vectors

import scala.collection.mutable

object VecUtil {


  /**
    * 一个udf函数，它可以输入一篇文档，输出一个tf值的特征向量（数组）
    */
  val doc2tf = (doc:String,n:Int)=>{

    //准备一个长度为n，初始值为0的数组 [0,0,0,0,0,0,0,0,0,0,.....]
    val tfArr = Array.fill(n)(0.0)


    // 字符串："a a a a a a x x y" => 字符串数组：[a,a,a,a,a,a,x,x,y]
    val tmp1: Array[String] = doc.split(" ")

    //数组：[a,a,a,a,a,a,x,x,y] =>  HashMap：{(a,[a,a,a,a,a,a]),(x,[x,x]),(y,[y])}
    val tmp2: Map[String, Array[String]] = tmp1.groupBy(e=>e)

    // (a,[a,a,a,a,a,a]) => (a,6)
    // (x,[x,x])         => (x,2)
    // (y,[y])           => (y,1)
    val wc: Map[String, Int] = tmp2.map(tp=>(tp._1,tp._2.size))

    // 遍历wc这个hashmap，将其中的每一个词的词频映射到向量的这个位置： 词.hashcode%n
    for((w,c)<-wc){
      // 用hash映射求得该词所映射的特征位置脚标
      val index = (w.hashCode&Integer.MAX_VALUE)%n
      // 再把特征向量中该脚标上的特征值，替换为这个词w的词频c
      tfArr(index) = c

    }

    tfArr
  }


  /**
    * 一个udf函数，它可以输入一个double数组，然后将数组中非0值替换成1，返回
    */
  val arr2One = (arr:mutable.WrappedArray[Double])=>{
    arr.map(d=>if(d != 0.0) 1.0 else 0.0)
  }


  /**
    * 将一个“词文档数” 数组 ==》 词idf值数组
    */
  val docCntArr2Idf = (doc_cnt:mutable.WrappedArray[Double],doc_total:Long)=>{

    doc_cnt.map(d=> Math.log10(doc_total/(0.01+d)))
  }


  /**
    * 将一个double数组，转为vector
    */
  val arr2Vec = (arr:mutable.WrappedArray[Double])=> Vectors.dense(arr.toArray)


}
